from docx import Document
from app.api.schemas import Lang, FragmentTransResultDict
from app.utils import TextSplitterUtil

# 参考：https://www.osgeo.cn/python-tutorial/docx-docx-tmpl.html
# 对于docx文档，可以将其转换为xml格式，然后将xml格式放到gpt中进行提问会比较方便和准确


class DocxUtils:
    def __init__(self, doc_path, from_lang: Lang, to_lang: Lang):
        self.doc = Document(doc_path)
        self.from_lang = from_lang
        self.to_lang = to_lang
        self.namespace = {
            "w": "http://schemas.openxmlformats.org/wordprocessingml/2006/main",
            "v": "urn:schemas-microsoft-com:vml",
            "o": "urn:schemas-microsoft-com:office:office",
            "w10": "urn:schemas-microsoft-com:office:word",
        }
        self.splitter = TextSplitterUtil()

    def extranct_text(self, trans_dict: FragmentTransResultDict):
        """提取待翻译文本"""
        for text in self._iter_texts():
            self._update_trans_dict(text, trans_dict)

    def replace_text(self, trans_dict: FragmentTransResultDict):
        """替换文本"""
        # 遍历文档中的所有段落
        for para in self.doc.paragraphs:
            if para.text.strip():
                self._replace_para(
                    para.runs, self._get_trans_target_text(para.text, trans_dict)
                )

        # 遍历文档中的所有表格
        for table in self.doc.tables:
            for row in table.rows:
                for cell in row.cells:
                    self._replace_cell(
                        cell, self._get_trans_target_text(cell.text, trans_dict)
                    )  # 替换表格中的文本

        # 遍历文档中的所有文本框
        root = self.doc.element  # 获取文档的 XML 根元素
        for txbxContent in root.findall(".//w:txbxContent", self.namespace):
            for paragraph in txbxContent.findall(".//w:p", self.namespace):
                for text_elem in paragraph.findall(".//w:t", self.namespace):
                    text_elem.text = self._get_trans_target_text(
                        text_elem.text, trans_dict
                    )  # 替换文本框中的文本

        # 遍历文档中的所有艺术字体文本框
        for pict in root.findall(".//w:pict", self.namespace):
            for shape in pict.findall(".//v:shape", self.namespace):
                textpath = shape.find(".//v:textpath", self.namespace)
                if textpath is not None and "string" in textpath.attrib:
                    textpath.attrib["string"] = self._get_trans_target_text(
                        textpath.attrib["string"], trans_dict
                    )  # 替换艺术字体文本框中的文本

    def save(self, output_path):
        """保存文档"""
        self.doc.save(output_path)

    def _iter_texts(self):
        """生成器函数，遍历文档中的所有文本"""
        # 遍历文档中的所有段落
        for para in self.doc.paragraphs:
            yield para.text

        # 遍历文档中的所有表格
        for table in self.doc.tables:
            for row in table.rows:
                for cell in row.cells:
                    yield cell.text

        # 遍历文档中的所有文本框
        root = self.doc.element  # 获取文档的 XML 根元素
        for txbxContent in root.findall(".//w:txbxContent", self.namespace):
            for paragraph in txbxContent.findall(".//w:p", self.namespace):
                for text_elem in paragraph.findall(".//w:t", self.namespace):
                    yield text_elem.text

        # 遍历文档中的所有艺术字体文本框
        for pict in root.findall(".//w:pict", self.namespace):
            for shape in pict.findall(".//v:shape", self.namespace):
                textpath = shape.find(".//v:textpath", self.namespace)
                if textpath is not None and "string" in textpath.attrib:
                    yield textpath.attrib["string"]

    def _update_trans_dict(self, text: str, trans_dict: FragmentTransResultDict):
        """更新翻译字典"""
        if text.strip() == "":
            return

        trans_dict.add_entity_from_details(text, self.from_lang, self.to_lang)

    def _get_trans_target_text(
        self, src_text: str, trans_dict: FragmentTransResultDict
    ):
        """获取翻译目标文本"""
        result = trans_dict.get_entity_by_src(src_text, self.to_lang)
        return result.dst if result else ""

    def _replace_para(self, paras, text):
        """
        替换段落中的文本，同时保留原有的样式
        """
        if not text:
            return

        chunks = self.splitter.split(text)
        run_count = len(paras)
        chunk_count = len(chunks)

        pos = 0
        for i in range(run_count):
            if pos >= chunk_count:
                paras[i].text = ""  # 清空多余的 run
                continue

            if i == run_count - 1:
                # 如果是最后一个 run，将 text 剩余全部给这个 run
                paras[i].text = "".join(chunks[pos:])
                break

            if paras[i].text.strip() == "":
                continue

            if pos <= chunk_count - 1:
                # 按 run 的文本长度，从 text 中截取
                paras[i].text = chunks[pos]
                pos += 1
                continue

            # 如果不够，则将剩余全部给 run
            paras[i].text = "".join(chunks[pos:])
            pos = chunk_count

    def _replace_cell(self, cell, text):
        """替换表格单元格中的文本，同时保留原有的样式"""
        self._replace_para(cell.paragraphs[0].runs, text)
